"""
LangChain 对话历史与会话隔离学习示例
===================================================
这个示例演示如何为Agent添加记忆功能，并实现多用户会话隔离。
关键概念：每个用户有独立的对话历史，互不干扰。

核心组件：
1. BaseChatMessageHistory: 聊天历史基类
2. InMemoryChatMessageHistory: 内存中的聊天历史实现
3. RunnableWithMessageHistory: 为Agent添加历史管理功能的包装器
4. Session ID: 会话标识符，用于区分不同用户

学习目标：
- 理解对话历史的管理机制
- 掌握会话隔离的实现方法
- 学会使用RunnableWithMessageHistory包装器
- 了解多用户场景下的对话管理
"""
import os
import time
from dotenv import load_dotenv
from langchain import hub
from langchain.agents import create_tool_calling_agent, AgentExecutor
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import FAISS
from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_core.tools import create_retriever_tool
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 自动加载 .env 文件中的环境变量（用于API密钥等配置）
load_dotenv()

# ==================== 定义大语言模型 ====================
llm = ChatOpenAI(api_key=os.getenv("MODELSCOPE_API_KEY"), openai_api_base=os.getenv("MODELSCOPE_API_BASE"), model="Qwen/Qwen2.5-VL-72B-Instruct", streaming=True, temperature=0.7)


# ==================== 知识库构建流程 ====================

# 步骤1：文档加载 - 从网页加载知识内容
print("正在加载网页知识内容...")
loader = WebBaseLoader("https://baike.baidu.com/item/%E4%BA%BA%E5%B7%A5%E6%99%BA%E8%83%BD/9180")
# 加载文档，返回Document对象列表，每个Document包含页面内容和元数据
docs = loader.load()
print(f"成功加载 {len(docs)} 个文档")

# 步骤2：文档切分 - 将长文档分割为适合处理的文本块
print("正在分割文档...")
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,  # 每个文本块的最大字符数
    chunk_overlap=200  # 块之间的重叠字符数，保持上下文连贯性
)
documents = text_splitter.split_documents(docs)  # 执行切分操作
print(f"文档分割完成，共 {len(documents)} 个文本块")

# 步骤3：向量化处理 - 将文本转换为数值向量
print("正在初始化向量化模型...")
embeddings = HuggingFaceEmbeddings(model_name="D:/models/BAAIbge-base-zh-v1.5", model_kwargs={'device': 'cpu'})


# 步骤4：创建向量数据库 - 存储文档向量用于相似度检索
print("正在构建向量数据库...")
vector = FAISS.from_documents(documents, embeddings)
print("向量数据库构建完成")

# 步骤5：创建检索器 - 用于在向量数据库中搜索相关文档
retriever = vector.as_retriever()

# ==================== 工具创建 ====================

# 创建检索工具 - 让Agent能够搜索百度百科知识库
retriever_tool = create_retriever_tool(
    retriever,  # 检索器对象
    "baike_search",  # 工具名称（Agent调用时使用的标识）
    "搜索百度百科中关于人工智能的相关知识"  # 工具描述，Agent根据描述决定何时使用该工具
)

# 工具列表 - 可以添加多个工具，Agent会根据问题自动选择
tools = [retriever_tool]
print(f"已创建 {len(tools)} 个工具")

# ==================== 代理系统构建 ====================

# 获取预定义的提示词模板
print("正在加载提示词模板...")
prompt = hub.pull("hwchase17/openai-functions-agent")  # 从LangChain Hub拉取标准模板

# 创建工具调用代理
print("正在创建智能代理...")
agent = create_tool_calling_agent(
    llm,  # 大语言模型，负责理解和决策
    tools,  # 可用工具列表
    prompt  # 提示词模板，指导Agent的行为
)

# 创建代理执行器 - 管理Agent的执行流程
agent_executor = AgentExecutor(
    agent=agent,  # 代理实例
    tools=tools,  # 工具列表
    verbose=False  # 不显示详细执行过程（测试时可设为True）
)

# ==================== 对话历史管理部分 ====================

"""
会话隔离核心概念：
------------------------------------
在真实应用场景中，我们通常需要同时服务多个用户。
每个用户应该有自己独立的对话历史，互不干扰。

实现方式：
1. 为每个用户分配唯一的 session_id
2. 使用字典存储不同session_id对应的聊天历史
3. 通过RunnableWithMessageHistory自动管理历史记录
"""

# 定义存储session的容器
# 使用字典来存储不同会话的聊天历史，key=session_id, value=聊天历史对象
store = {}


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    """
    根据session_id获取或创建聊天历史的工厂函数

    参数:
        session_id: 会话标识符，用于区分不同用户的对话
                   例如："user_123", "customer_456", "niuniu"等

    返回:
        BaseChatMessageHistory: 聊天历史对象，存储该会话的所有消息

    工作机制：
        - 如果session_id已存在，返回对应的历史记录
        - 如果session_id不存在，创建新的内存聊天历史并存储
    """
    # 如果session_id不存在于store中，创建一个新的内存聊天历史
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
        print(f"创建新的会话历史: {session_id}")
    else:
        print(f"使用现有会话历史: {session_id}")

    return store[session_id]


# 创建带有聊天历史功能的Agent包装器
# RunnableWithMessageHistory 为Agent添加了记忆能力，使其能够记住对话上下文
agent_with_chat_history = RunnableWithMessageHistory(
    runnable=agent_executor,  # 基础的Agent执行器（无记忆功能）
    get_session_history=get_session_history,  # 获取历史记录的工厂函数
    input_messages_key="input",  # 输入消息在输入字典中的键名
    history_messages_key="chat_history"  # 历史记录在输入字典中的键名

    # 其他可选参数：
    # output_messages_key: 输出消息的键名（默认"output"）
    # history_factory_config: 历史记录工厂的配置
)

print("✓ 带历史记录的Agent包装器创建完成")

# ==================== 测试对话部分 ====================

print("\n" + "=" * 60)
print("测试1：用户'牛牛'的对话")
print("=" * 60)

# 第一次对话：用户自我介绍
print("\n--- 第一次对话：自我介绍 ---")
response1 = agent_with_chat_history.invoke(
    {"input": "Hi,我的名字叫牛牛"},  # 用户输入
    config={"configurable": {"session_id": "niuniu"}}  # 配置session_id为"niuniu"
)
print("AI回复:", response1['output'])
print("当前会话历史:", [msg.type for msg in store["niuniu"].messages])

# 添加请求间隔，模拟真实对话场景
time.sleep(2)

# 第二次对话：测试记忆功能 - 询问用户名字
print("\n--- 第二次对话：测试记忆功能 ---")
response2 = agent_with_chat_history.invoke(
    {"input": "我叫什么名字"},  # 测试Agent是否能记住之前的对话
    config={"configurable": {"session_id": "niuniu"}}  # 使用相同的session_id
)
print("AI回复:", response2['output'])
print("当前会话历史:", [msg.type for msg in store["niuniu"].messages])

# ==================== 多用户会话隔离测试 ====================

print("\n" + "=" * 60)
print("测试2：用户'小花'的对话（会话隔离测试）")
print("=" * 60)

# 新用户对话：测试会话隔离
print("\n--- 新用户'小花'的对话 ---")
response3 = agent_with_chat_history.invoke(
    {"input": "你好，我是小花"},  # 新用户输入
    config={"configurable": {"session_id": "xiaohua"}}  # 不同的session_id
)
print("AI回复:", response3['output'])
print("小花会话历史:", [msg.type for msg in store["xiaohua"].messages])

# 测试会话隔离：询问小花的名字
print("\n--- 测试会话隔离 ---")
response4 = agent_with_chat_history.invoke(
    {"input": "我的名字是什么"},  # 应该正确回答"小花"
    config={"configurable": {"session_id": "xiaohua"}}  # 继续使用小花的session_id
)
print("AI回复:", response4['output'])

# 验证隔离：用牛牛的session_id询问名字
print("\n--- 验证会话隔离 ---")
response5 = agent_with_chat_history.invoke(
    {"input": "我的名字是什么"},  # 应该正确回答"牛牛"
    config={"configurable": {"session_id": "niuniu"}}  # 切换回牛牛的session_id
)
print("AI回复:", response5['output'])

# ==================== 系统状态展示 ====================

print("\n" + "=" * 60)
print("系统状态总结")
print("=" * 60)
print(f"当前管理的会话数量: {len(store)}")
print(f"活跃会话ID: {list(store.keys())}")

for session_id, history in store.items():
    print(f"\n会话 '{session_id}' 的消息统计:")
    print(f"  - 总消息数: {len(history.messages)}")
    human_msgs = [msg for msg in history.messages if msg.type == 'human']
    ai_msgs = [msg for msg in history.messages if msg.type == 'ai']
    print(f"  - 用户消息: {len(human_msgs)} 条")
    print(f"  - AI回复: {len(ai_msgs)} 条")

"""
扩展应用场景：
------------------------------------
1. 多用户聊天机器人：每个用户有独立的对话历史
2. 客服系统：每个客户会话保持上下文
3. 个性化助手：基于用户历史提供个性化服务
4. 教学系统：每个学生有独立的学习记录

持久化存储方案：
------------------------------------
当前使用内存存储，重启后数据会丢失。生产环境可以考虑：
1. 数据库存储：MySQL, PostgreSQL
2. 缓存系统：Redis, Memcached  
3. 文件存储：JSON文件
4. 云服务：Firestore, DynamoDB

示例：使用Redis持久化
from langchain_community.chat_message_histories import RedisChatMessageHistory

def get_redis_session_history(session_id: str):
    return RedisChatMessageHistory(session_id, url="redis://localhost:6379")
"""

print("\n" + "=" * 60)
print("学习要点总结")
print("=" * 60)
print("✓ 掌握了对话历史管理的基本概念")
print("✓ 学会了使用RunnableWithMessageHistory包装器")
print("✓ 理解了会话隔离的实现原理")
print("✓ 实践了多用户场景下的对话管理")
print("✓ 了解了session_id的作用和配置方法")

print("\n核心优势：")
print("1. 自动历史管理：无需手动维护聊天历史")
print("2. 会话隔离：不同用户对话完全独立")
print("3. 灵活扩展：支持多种存储后端")
print("4. 无缝集成：与现有Agent系统完美结合")